package com.md.client.helper;

import com.md.context.LogServerContext;
import org.apache.http.client.methods.HttpRequestBase;
import org.aspectj.lang.ProceedingJoinPoint;
import org.aspectj.lang.annotation.Around;
import org.aspectj.lang.annotation.Aspect;
import org.aspectj.lang.annotation.Pointcut;
import org.springframework.http.HttpEntity;
import org.springframework.http.HttpHeaders;
import org.springframework.stereotype.Component;

import static com.md.value.LogsConstants.SERVER_CONTEXT_KEY_DATA_ID;
import static com.md.value.LogsConstants.SERVER_CONTEXT_KEY_IP_ADDRESS;

/**
 * @author 刘骄阳
 * 2024-10-16 18:54
 */
@Aspect
@Component
public class HttpHeaderInterceptorAspect {

  @Pointcut("execution(* org.springframework.web.client.RestTemplate.exchange(..))")
  public void restTemplateExchange() {
  }

  @Pointcut("execution(* org.apache.http.client.methods.HttpRequestBase..*(..))")
  public void httpClientRequest() {
  }

  @Pointcut("execution(* okhttp3.Request$Builder.header(..))")
  public void okHttpRequest() {
  }

  @Around("restTemplateExchange()")
  public Object aroundRestTemplateExchange(final ProceedingJoinPoint joinPoint) throws Throwable {
    Object[] args = joinPoint.getArgs();
    for (int i = 0; i < args.length; i++) {
      Object arg = args[i];
      if (arg instanceof HttpEntity) {
        HttpEntity<?> httpEntity = (HttpEntity<?>) arg;
        HttpEntity<?> newEntity = new HttpEntity<>(httpEntity.getBody(), handleRestTemplate(httpEntity.getHeaders()));
        args[i] = newEntity;
      }
    }
    // 继续执行目标方法，并返回结果
    return joinPoint.proceed(args);
  }

  @Around("httpClientRequest()")
  public Object addTraceIdHeaderForHttpClient(final ProceedingJoinPoint joinPoint) throws Throwable {
    Object[] args = joinPoint.getArgs();
    for (Object arg : args) {
      if (arg instanceof HttpRequestBase) {
        handleHttpClient((HttpRequestBase) arg);
      }
    }  // 继续执行目标方法，并返回结果
    return joinPoint.proceed(args);
  }

  @Around("okHttpRequest()")
  public Object addTraceIdHeaderForOkHttp(final ProceedingJoinPoint joinPoint) throws Throwable {
    Object[] args = joinPoint.getArgs();
    for (Object arg : args) {
      if (arg instanceof okhttp3.Request.Builder) {
        handleOkHttp((okhttp3.Request.Builder) arg);
      }
    } // 继续执行目标方法，并返回结果
    return joinPoint.proceed(args);
  }

  private HttpHeaders handleRestTemplate(final HttpHeaders httpHeaders) {
    HttpHeaders headers = new HttpHeaders();
    // 复制原有头部
    headers.putAll(httpHeaders);
    headers.add(SERVER_CONTEXT_KEY_DATA_ID, LogServerContext.getString(SERVER_CONTEXT_KEY_DATA_ID));
    headers.add(SERVER_CONTEXT_KEY_IP_ADDRESS, LogServerContext.getString(SERVER_CONTEXT_KEY_IP_ADDRESS));
    return headers;
  }

  private void handleHttpClient(final HttpRequestBase httpRequestBase) {
    httpRequestBase.addHeader(SERVER_CONTEXT_KEY_DATA_ID, LogServerContext.getString(SERVER_CONTEXT_KEY_DATA_ID));
    httpRequestBase.addHeader(SERVER_CONTEXT_KEY_IP_ADDRESS, LogServerContext.getString(SERVER_CONTEXT_KEY_IP_ADDRESS));
  }

  private void handleOkHttp(final okhttp3.Request.Builder builder) {
    builder.header(SERVER_CONTEXT_KEY_DATA_ID, LogServerContext.getString(SERVER_CONTEXT_KEY_DATA_ID));
    builder.header(SERVER_CONTEXT_KEY_IP_ADDRESS, LogServerContext.getString(SERVER_CONTEXT_KEY_IP_ADDRESS));
  }
}
